import torch
import time
from ptflops import get_model_complexity_info

# 假设你已经有一个训练好的模型类
# 这里使用一个简单的模型架构举例，你可以用你自己的模型替代
import torchvision.models as models

# 1. 加载已经训练好的模型
# 假设模型保存在 model.pth 文件中
model = models.resnet50(pretrained=True)  # 用你的模型替代

# 加载训练好的权重
model.load_state_dict(torch.load('co_deformable_detr_r50_1x_coco.pth'), strict=False)
model.eval()  # 切换到评估模式
def get_flops_pytorch(model, batch):
    # 处理模型的 forward 过程
    # 确保 batch 是 4D 张量
    batch = batch.view(batch.size(0), -1, batch.size(2), batch.size(3))  # 对 5D 输入做变换
    # 继续其余的计算

# 2. 计算模型的参数量和 GFLOPs（通过 ptflops）
def calculate_complexity(model, input_size=(1, 3, 224, 224)):
    with torch.cuda.device(0):  # 如果你有 GPU，确保模型和计算在 GPU 上进行
        macs, params = get_model_complexity_info(model, input_size, as_strings=True, print_per_layer_stat=True)
        print(f"Model Parameters: {params}")
        print(f"GFLOPs: {macs}")

# 计算参数量和 GFLOPs
calculate_complexity(model)

# 3. 计算推理速度 (FPS)
def calculate_inference_speed(model, input_size=(1, 3, 224, 224), iterations=100):
    model.eval()  # 切换到评估模式
    input_data = torch.randn(*input_size).to('cuda')  # 生成随机输入数据
    model = model.to('cuda')  # 将模型转移到 GPU

    start_time = time.time()
    for _ in range(iterations):
        with torch.no_grad():
            output = model(input_data)  # 执行一次推理

    total_time = time.time() - start_time
    fps = iterations / total_time  # FPS 计算
    return fps

# 计算推理速度 (FPS)
fps = calculate_inference_speed(model)
print(f"Estimated FPS: {fps} FPS")
